import xarray as xr
import numpy as np
import rasterio
from rasterio.transform import from_origin
import os


def nc_to_tif(input_nc_path, variable_name, output_tif_path=None):
    """
    将 NetCDF 文件转换为 GeoTIFF 文件，并计算全年累计暴雨量。

    参数:
    - input_nc_path: 输入的 NetCDF 文件路径。
    - variable_name: NetCDF 文件中的变量名。
    - output_tif_path: 输出的 GeoTIFF 文件路径。如果未指定，则自动生成。
    """
    # 读取 NetCDF 数据
    precip_dataset = xr.open_dataset(input_nc_path)

    # 提取指定变量的数据
    if variable_name not in precip_dataset.data_vars:
        raise ValueError(f"变量 '{variable_name}' 不存在于 NetCDF 文件中。")

    tp = precip_dataset[variable_name]

    # todo 开发测试时使用
    # latitude_range = slice(55, 54)  # 纬度范围
    # longitude_range = slice(70, 71)  # 经度范围
    #
    # # 指定时间范围
    # start_time = '2023-01-01T00:00:00'  # 2023年1月1日12:00
    # end_time = '2023-01-05T23:00:00'  # 2023年1月2日12:00
    #
    # tp_selected = tp.sel(latitude=latitude_range, longitude=longitude_range, valid_time=slice(start_time, end_time))

    tp_m = tp

    # 计算每天的总降雨量
    tp_daily = tp_m.resample(valid_time='1D').sum(dim='valid_time')

    # 初始化全年累计暴雨量
    # R_storm = np.zeros_like(tp_daily.isel(valid_time=0).values)
    R_storm = np.zeros_like(tp_daily.values)

    # 暴雨标准判断 若24小时降雨量≥50 mm，则该天的降雨量计入全年累计暴雨量。
    for i in range(len(tp_daily.valid_time)):
        daily_rainfall = tp_daily.isel(valid_time=i).values
        storm_mask = daily_rainfall >= 50.0 / 1000    # 将每小时降雨量从毫米转换为米
        R_storm += daily_rainfall * storm_mask

    # 需要对R_storm去掉1维数据，然后做汇总
    # 去掉多余的维度
    R_storm = R_storm.squeeze()
    R_storm = R_storm.sum(axis=0)
    # 设置输出 GeoTIFF 的参数
    transform = from_origin(
        tp.longitude[0].item(),
        tp.latitude[0].item(),
        tp.longitude[1].item() - tp.longitude[0].item(),
        tp.latitude[0].item() - tp.latitude[1].item()
    )

    # 设置 TIFF 文件的元数据
    metadata = {
        'driver': 'GTiff',
        'count': 1,
        'dtype': R_storm.dtype,
        'width': R_storm.shape[1],
        'height': R_storm.shape[0],
        'crs': 'EPSG:4326',
        'transform': transform
    }

    # 如果未指定输出路径，则自动生成
    if output_tif_path is None:
        base_name = os.path.basename(input_nc_path).replace('.nc', '')
        output_tif_path = os.path.join(os.path.dirname(input_nc_path), f"{base_name}_heavy_precip_1YE.tif")

    # 写入 TIFF 文件
    with rasterio.open(output_tif_path, 'w', **metadata) as dst:
        dst.write(R_storm, 1)  # 写入第一个波段

        # 显式生成统计信息（关键修复）
        print(f"正在写入统计信息...")
        dst.write_colormap(1, {})  # 清空伪彩色表
        dst.update_tags(ns='rio_stats', STATISTICS_VALID_PERCENT=100)  # 标记有效统计
        # 强制重新计算统计信息（排除无效值）
        # stats = dst.stats(1)
        stats = dst.statistics(1)  # 或者使用正确的属性名
        print(f"实际数据范围: {stats.min} ~ {stats.max}")  # 验证输出范围

    print(f"GeoTIFF 文件已保存到: {output_tif_path}")
    return output_tif_path


# 示例调用
input_nc_path = 'G:/GEP_data/input/era5/2023_precip_hour.nc'
variable_name = 'tp'
output_tif_path = nc_to_tif(input_nc_path, variable_name)